import numpy as np
import matplotlib.pyplot as plt
from scipy import signal

# 1. 定义系统传递函数 H(s) = 1/(s + 1)
num = [1]     # 分子系数: 1
den = [1, 1]  # 分母系数: s + 1
sys = signal.TransferFunction(num, den)

# 2. 定义输入信号 f2(t)（严格匹配用户图片）
t = np.linspace(0, 3, 1000)
f2 = np.piecewise(
    t,
    [t < 1, t >= 1],
    [
        lambda t: (3/2) * t,             # 0 ≤ t < 1: 线性上升到 1.5
        lambda t: (3/2) * (2 - (t - 1)) # 1 ≤ t ≤ 3: 线性下降到 0
    ]
)

# 3. 计算零状态响应（使用 scipy.signal.lsim）
t, y, _ = signal.lsim(sys, U=f2, T=t)

# 4. 绘制输入信号和响应曲线
plt.figure(figsize=(10, 6))

# 输入信号 f2(t)
plt.subplot(2, 1, 1)
plt.plot(t, f2, label='Input: $f_2(t)$', color='purple')
plt.fill_between(t, f2, color='purple', alpha=0.3)  # 填充紫色区域
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.xticks([0, 1, 2, 3])
plt.yticks([0, 1.5])
plt.grid(True, linestyle='--', alpha=0.6)
plt.legend()

# 零状态响应
plt.subplot(2, 1, 2)
plt.plot(t, y, label='Zero-State Response', color='red')
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.xticks([0, 1, 2, 3])
plt.grid(True, linestyle='--', alpha=0.6)
plt.legend()

plt.tight_layout()
plt.show()